from flask import Blueprint, jsonify,request
import pandas as pd
import os
import json
from datetime import datetime
from lib.lib import get_price

real_stock_data = Blueprint('real_stock_data', __name__)

@real_stock_data.route('/pullAllRealStockData', methods=['GET', 'POST', 'DELETE', 'PUT'])
def pullAllRealStockData():
    type = request.args.get('type')
    try:
        # 从stock_info.json文件读取股票列表
        stock_info_path = 'data/stock_info.json'
        
        if not os.path.exists(stock_info_path):
            return jsonify({
                'code': 404,
                'message': 'stock_info.json file not found'
            }), 404
        
        with open(stock_info_path, 'r', encoding='utf-8') as f:
            stock_data = json.load(f)
        
        # 提取股票代码列表
        stocks = [stock['code'] for stock in stock_data]
        
        for code in stocks:
            if type == 'history':
                csv_path = f'data/5m/{code}.csv'
            else:
                csv_path = f'data/realgo_5m/{code}.csv'
            
            # 确保目录存在
            os.makedirs(os.path.dirname(csv_path), exist_ok=True)
            
            # 读取现有数据（如果存在）
            try:
                existing_df = pd.read_csv(csv_path)
                existing_df['date'] = pd.to_datetime(existing_df['date'])
                last_date = existing_df['date'].max()
                print(f"Last date for {code}: {last_date}")
            except FileNotFoundError:
                existing_df = pd.DataFrame()
                last_date = pd.Timestamp('2013-12-02')  # 默认起始日期
                print(f"No existing data for {code}, starting from {last_date}")
            
            # 获取新数据
            try:
                new_df = get_price(code, frequency='5m', count=100)
                if len(new_df) == 0:
                    print(f"No new data available for {code}")
                    continue
                
                # 处理新数据
                new_df = new_df.reset_index()
                new_df.columns = ['date'] + list(new_df.columns[1:])
                new_df['code'] = code
                
                # 如果有现有数据，只保留更新的部分
                if not existing_df.empty:
                    new_df = new_df[new_df['date'] > last_date]
                
                if len(new_df) == 0:
                    print(f"No new data to append for {code}")
                    continue
                
                # 合并数据
                if existing_df.empty:
                    final_df = new_df
                else:
                    final_df = pd.concat([existing_df, new_df], ignore_index=True)
                
                # 只保留需要的列
                final_df = final_df[['date', 'close', 'code', 'open', 'high', 'low', 'volume']]
                
                # 按日期排序
                final_df = final_df.sort_values('date')
                
                # 去除重复数据
                final_df = final_df.drop_duplicates(subset=['date'], keep='last')
                
                # 保存到CSV
                final_df.to_csv(csv_path, index=False)
                print(f"Successfully updated {code}, added {len(new_df)} new records")
                
            except Exception as e:
                print(f"Error processing {code}: {str(e)}")
                continue
        
        return jsonify({
            'code': 200,
            'message': 'Data update completed'
        })

    except Exception as e:
        print(f"Error in pullAllRealStockData: {str(e)}")
        return jsonify({
            'code': 500,
            'message': f'Error: {str(e)}'
        }), 500 